// Script to read out of events.lhe from MadGraph and create a ROOT file events.lhe.root

using namespace std;

#include <vector>
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <cstring>
#include <TFile.h>
#include <TTree.h>
#include <TCanvas.h>
#include <TFrame.h>
#include <TH1F.h>
#include <TBenchmark.h>
#include <TRandom.h>
#include <TSystem.h>

void treeCreationLHCO(char* fileName) 
{
  
  // Read input file and skip up to where data starts
  string line;
  const char* source = fileName;
  ifstream myfile(source);
  cout << "Reading from: " << source << endl;
  while(getline(myfile,line)) {
    if(line=="#</LesHouchesEvents>") {
      getline(myfile,line);
      getline(myfile,line);
      getline(myfile,line);
      getline(myfile,line);
      break;
    }
  }

  // Open .root file
  char filename[100]="";
  const char* fileExtension = ".root";
  strcat(filename,source);
  strcat(filename,fileExtension);
  TFile *f = TFile::Open(filename,"RECREATE");   
  if (!f) { return; }
  cout << "Saving to: " << filename << endl;

  Int_t run = 1989;

  Int_t event = 0;
  Int_t nentries = 0;
  Int_t nNoMuonEvents = 0; //events without muons
  Int_t nmuons = 0; //muons per event

  std::vector<int> typ;
  std::vector<float> eta;
  std::vector<float> phi;
  std::vector<float> pt;
  std::vector<float> jmas;
  std::vector<float> ntrk;
  std::vector<float> btag;
  std::vector<float> hadEm;

  // Create a TTree
  TTree *t = new TTree("tree","Tree with vectors");
  t->Branch("run",&run);
  t->Branch("event",&event);
  t->Branch("nentries",&nentries);
  t->Branch("typ",&typ);
  t->Branch("eta",&eta);
  t->Branch("phi",&phi);
  t->Branch("pt",&pt);
  t->Branch("jmas",&jmas);
  t->Branch("ntrk",&ntrk);
  t->Branch("btag",&btag);
  t->Branch("hadEm",&hadEm);

  Int_t nevents = 0;
  while(getline(myfile,line)) {
    stringstream ss(line);
    Int_t firstEntry;
    ss >> firstEntry;
    if(firstEntry == 0) {
      if(nentries!=0) {
	event = nevents;
	nevents++;
	if(nmuons==0) {
	  nNoMuonEvents++;
	}
	else {
	  t->Fill();
	}
      }

      
      typ.clear();
      eta.clear();
      phi.clear();
      pt.clear();
      jmas.clear();
      ntrk.clear();
      btag.clear();
      hadEm.clear();
      nmuons = 0;
      //getline(myfile,line);
    }
    else {
      Float_t leta, lphi, lpt, ljmas, lntrk, lbtag, lhadEm;
      Int_t ltyp;
      ss >> ltyp >> leta >> lphi >> lpt >> ljmas >> lntrk >> lbtag >> lhadEm;
      typ.push_back(ltyp);
      eta.push_back(leta);
      phi.push_back(lphi);
      pt.push_back(lpt);
      jmas.push_back(ljmas);
      ntrk.push_back(lntrk);
      btag.push_back(lbtag);
      hadEm.push_back(lhadEm);
      nentries++;
      if(ltyp==2) {
	nmuons++;
      }
    }

  } 

  f->Write();
  myfile.close();
  cout << "Number of events read in:" << nevents << endl;
  cout << nNoMuonEvents << " events without muons not thrown away" << endl;
  cout << "Final number of events in histogram: " << nevents-nNoMuonEvents << endl;
  cout << filename << " created!" << endl;
  delete f;
}

int main(int argc, char **argv) {
  std::cout << "Tree creation LHCO started!" << std::endl;
  treeCreationLHCO(argv[1]);
  std::cout << "Tree creation LHCO done!" << std::endl;
}
